import json
import time
from dataclasses import dataclass
from datetime import datetime
from typing import Optional, List, Dict, Any

import requests
import zmq


@dataclass
class WorkflowRun:
    workflow_run_id: str
    workflow_id: str
    user_id: str
    status: str
    started_at: datetime
    finished_at: Optional[datetime] = None
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

    @classmethod
    def from_dict(cls, data: dict) -> 'WorkflowRun':
        return cls(
            workflow_run_id=data['workflow_run_id'],
            workflow_id=data['workflow_id'],
            user_id=data['user_id'],
            status=data['status'],
            started_at=datetime.fromisoformat(data['started_at'].replace('Z', '+00:00')),
            finished_at=datetime.fromisoformat(data['finished_at'].replace('Z', '+00:00')) if data.get('finished_at') else None,
            created_at=datetime.fromisoformat(data['created_at'].replace('Z', '+00:00')) if data.get('created_at') else None,
            updated_at=datetime.fromisoformat(data['updated_at'].replace('Z', '+00:00')) if data.get('updated_at') else None,
        )


@dataclass
class Trace:
    event: str
    workflow_run_id: str
    workflow_id: str
    node_run_id: str
    node_id: str
    type: int
    data: Dict[str, Any]
    created_at: Optional[datetime] = None

    @classmethod
    def from_dict(cls, data: dict) -> 'Trace':
        return cls(
            event=data['event'],
            workflow_run_id=data['workflow_run_id'],
            workflow_id=data['workflow_id'],
            node_run_id=data['node_run_id'],
            node_id=data['node_id'],
            type=data['type'],
            data=data['data'],
            created_at=datetime.fromisoformat(data['created_at'].replace('Z', '+00:00')) if data.get('created_at') else None,
        )


class AWMClient:
    def __init__(self, base_url: str, zmq_url: str, api_key: str):
        """
        初始化 AWM 客户端

        :param base_url: AWM 服务器基础URL，例如：http://localhost:8080/api/v1
        :param zmq_url: ZeroMQ 服务器URL，例如：tcp://localhost:5555
        :param api_key: API Key
        """
        self.base_url = base_url.rstrip('/')
        self.zmq_url = zmq_url
        self.api_key = api_key
        self.headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }

        # 初始化 ZeroMQ
        self.zmq_context = zmq.Context()
        self.zmq_socket = self.zmq_context.socket(zmq.PUSH)
        self.zmq_socket.connect(zmq_url)

    def create_workflow_run(self, workflow_run_id: str, workflow_id: str) -> WorkflowRun:
        """
        创建工作流运行记录

        :param workflow_run_id: 工作流运行ID
        :param workflow_id: 工作流ID
        :return: 工作流运行记录
        """
        url = f"{self.base_url}/workflow-runs"
        data = {
            "workflow_run_id": workflow_run_id,
            "workflow_id": workflow_id
        }
        response = requests.post(url, headers=self.headers, json=data)
        response.raise_for_status()
        return WorkflowRun.from_dict(response.json())

    def update_workflow_run(self, workflow_run_id: str, status: str, finished_at: Optional[datetime] = None) -> WorkflowRun:
        """
        更新工作流运行记录

        :param workflow_run_id: 工作流运行ID
        :param status: 状态（running/succeeded/failed）
        :param finished_at: 完成时间
        :return: 更新后的工作流运行记录
        """
        url = f"{self.base_url}/workflow-runs/{workflow_run_id}"
        data = {
            "status": status,
            "finished_at": finished_at.isoformat() if finished_at else None
        }
        response = requests.put(url, headers=self.headers, json=data)
        response.raise_for_status()
        return WorkflowRun.from_dict(response.json())

    def get_workflow_run(self, workflow_run_id: str) -> WorkflowRun:
        """
        获取工作流运行记录

        :param workflow_run_id: 工作流运行ID
        :return: 工作流运行记录
        """
        url = f"{self.base_url}/workflow-runs/{workflow_run_id}"
        response = requests.get(url, headers=self.headers)
        response.raise_for_status()
        return WorkflowRun.from_dict(response.json())

    def get_workflow_traces(self, workflow_run_id: str) -> List[Trace]:
        """
        获取工作流追踪日志

        :param workflow_run_id: 工作流运行ID
        :return: 追踪日志列表
        """
        url = f"{self.base_url}/traces/workflow-run/{workflow_run_id}"
        response = requests.get(url, headers=self.headers)
        response.raise_for_status()
        return [Trace.from_dict(item) for item in response.json()['items']]

    def send_trace(self, trace: Trace):
        """
        发送追踪日志

        :param trace: 追踪日志
        """
        data = {
            "event": trace.event,
            "workflow_run_id": trace.workflow_run_id,
            "workflow_id": trace.workflow_id,
            "node_run_id": trace.node_run_id,
            "node_id": trace.node_id,
            "type": trace.type,
            "data": trace.data
        }
        self.zmq_socket.send_json(data)

    def send_workflow_started(self, workflow_run_id: str, workflow_id: str, user_id: str):
        """
        发送工作流开始事件

        :param workflow_run_id: 工作流运行ID
        :param workflow_id: 工作流ID
        :param user_id: 用户ID
        """
        trace = Trace(
            event="workflow_started",
            workflow_run_id=workflow_run_id,
            workflow_id=workflow_id,
            node_run_id="",
            node_id="",
            type=0,
            data={
                "inputs": {
                    "sys.user_id": user_id,
                },
                "created_at": int(time.time())
            }
        )
        self.send_trace(trace)

    def send_workflow_finished(self, workflow_run_id: str, workflow_id: str, status: str):
        """
        发送工作流完成事件

        :param workflow_run_id: 工作流运行ID
        :param workflow_id: 工作流ID
        :param status: 状态（succeeded/failed）
        """
        trace = Trace(
            event="workflow_finished",
            workflow_run_id=workflow_run_id,
            workflow_id=workflow_id,
            node_run_id="",
            node_id="",
            type=0,
            data={
                "status": status,
                "finished_at": int(time.time())
            }
        )
        self.send_trace(trace)

    def send_node_started(self, workflow_run_id: str, workflow_id: str, node_run_id: str, node_id: str, node_type: int, data: dict):
        """
        发送节点开始事件

        :param workflow_run_id: 工作流运行ID
        :param workflow_id: 工作流ID
        :param node_run_id: 节点运行ID
        :param node_id: 节点ID
        :param node_type: 节点类型（0-工作流节点，1-组件节点）
        :param data: 节点数据
        """
        trace = Trace(
            event="node_started",
            workflow_run_id=workflow_run_id,
            workflow_id=workflow_id,
            node_run_id=node_run_id,
            node_id=node_id,
            type=node_type,
            data=data
        )
        self.send_trace(trace)

    def send_node_finished(self, workflow_run_id: str, workflow_id: str, node_run_id: str, node_id: str, node_type: int, data: dict):
        """
        发送节点完成事件

        :param workflow_run_id: 工作流运行ID
        :param workflow_id: 工作流ID
        :param node_run_id: 节点运行ID
        :param node_id: 节点ID
        :param node_type: 节点类型（0-工作流节点，1-组件节点）
        :param data: 节点数据
        """
        trace = Trace(
            event="node_finished",
            workflow_run_id=workflow_run_id,
            workflow_id=workflow_id,
            node_run_id=node_run_id,
            node_id=node_id,
            type=node_type,
            data=data
        )
        self.send_trace(trace)

    def close(self):
        """
        关闭客户端连接
        """
        self.zmq_socket.close()
        self.zmq_context.term() 